//Programmer: Aaron Steele
//Purpose: Compute the GC3 content for RRTs
//Usage: gc3 <fasta sequence_file> <codonmultipliers output>

#include <stdio.h>
#include <unistd.h>
#include <ctype.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <errno.h>
#include <stdbool.h>

int unique_codons = 64;
int unique_aa = 21;

/*An array of all possible codons*/
char codon_types[64][4] = {"TTT", "TTC", "TTA","TTG", "CTT", "CTC", "CTA", "CTG",
                         "ATT", "ATC", "ATA", "ATG", "GTT", "GTC", "GTA", "GTG",
                         "TCT", "TCC", "TCA", "TCG", "AGT", "AGC", "CCT", "CCC",
                         "CCA", "CCG", "ACT", "ACC", "ACA", "ACG", "GCT", "GCC",
                         "GCA", "GCG", "TAT", "TAC", "CAT", "CAC", "CAA", "CAG",
                         "AAT", "AAC", "AAA", "AAG", "GAT", "GAC", "GAA", "GAG",
                         "TGT", "TGC", "TGG", "CGT", "CGC", "CGA", "CGG", "AGA",
                         "AGG", "GGT", "GGC", "GGA", "GGG", "TAA", "TAG", "TGA"};

/* Maps each codon in codon_types to the corresponding amino acid */
char aa_codon_map[64] = {'F', 'F', 'L', 'L', 'L', 'L', 'L', 'L', 'I', 'I', 'I',
                         'M', 'V', 'V', 'V', 'V', 'S', 'S', 'S', 'S', 'S', 'S',
                         'P', 'P', 'P', 'P', 'T', 'T', 'T', 'T', 'A', 'A', 'A',
                         'A', 'Y', 'Y', 'H', 'H', 'Q', 'Q', 'N', 'N', 'K', 'K',
                         'D', 'D', 'E', 'E', 'C', 'C', 'W', 'R', 'R', 'R', 'R',
                         'R', 'R', 'G', 'G', 'G', 'G', 'X', 'X', 'X'};

/*An array of all amino acids (x corresponds to stop codon)*/
char aa_types[21] = {'F','L','I', 'M', 'V', 'S', 'P', 'T', 'A', 'Y', 'H', 'Q',
                     'N', 'K', 'D', 'E', 'C', 'W', 'R', 'G', 'X'};

/* Translate a codon to the corresponding Amino Acid */
char codon2aa(char *codon){
    int i;
    for(i=0; i<unique_codons; ++i){
        if(strcmp(codon,codon_types[i]) == 0)
            return aa_codon_map[i];
    }
}

/* Get the index of a given codon from the codon_types array */
int get_codon_index(char *codon){
    int i;
    for(i=0; i<unique_codons; ++i){
        if(strcmp(codon,codon_types[i]) == 0){
            return i;
        }
    }
    return -1;
}

/* Get the index of the amino acid with respect to the aa_types */
int get_aa_index(char aa){
    int i;
    for(i=0; i<unique_aa; ++i){
        if(aa == aa_types[i])
            return i;
    }
    return -1;
}

/*Get the next sequence from a given input fasta file*/
int getSequence(FILE *in, char seq[]){
    int d,s, length = 0;
    char* descr[300];

    if(in == NULL){
        return -1;
    }

    //Read the description of a sequence which ends with \n
    d = fscanf(in, "%[^\n]%*c", descr);
    
	if(d == EOF)
    	return -1;
    
    //Read the actual sequence which is terminated by the '>'
    s = fscanf(in, "%[^>]%n%*c", seq, &length);
    
	if(s == EOF)
		return -1;

	//Remove \n from the end of the sequence
    if(seq[length-1] == '\n'){
    	seq[length-1] = '\0';
    	length = length -1;
    }
    
    //Make all the characters uppercase for quality assurance
    int i;
    for(i=0; i<length; ++i)
    	seq[i] = toupper(seq[i]);
    
    return length;
}

//Check to see if given nucleotides are a valid codon
bool is_valid_codon(char* nucleotides){
	int index = get_codon_index(nucleotides);
	if(index == -1){
		return false;
	}

	return true;
}


void main(int argc, char **argv){
	
	if(argc !=3){
		fprintf(stderr, "No input/output fasta files specified, or too many arguments\n");
		exit(1);
	}

	//Input-Output File variables
	FILE *seqFile;
	FILE *mulFile;
	FILE *outFile;
	int index = 1;
	char *filename = NULL;

	//Bin and Seq Variables
	int binsize = 5;
	int bins = 100/binsize;
	int used_seq = 0;
	int bad_seq = 0;

	//Open the Input files
	filename = argv[index];
	seqFile = fopen(filename, "rt");
	if(seqFile == NULL){
		fprintf(stderr, "Failed to open %s\n", filename);
		exit(1);
	}

	char str[80];
	strcpy(str,filename);
	strcat(str,".codonsgc3");
	outFile = fopen(str, "wt");
	if(outFile == NULL){
		fprintf(stderr, "Failed to open %s\n", str);
		exit(1);
	}
	index++;

	filename = argv[index];
	mulFile = fopen(filename, "rt");
	if(mulFile == NULL){
		fprintf(stderr, "Failed to open %s\n", filename);
		exit(1);
	}
	
	//Arrays for counting GC content
	int seq_count[bins];
	double sum_gc3[bins];
	int codons[unique_codons];
	int codon_bins[bins][unique_codons];
	
	//Counting/Parsing variables
	int i,j,k;
	char seq[40000];
	int seq_length = -1;
	char cur_cdn[3];

	//Initialize
	for(i=0; i<bins; ++i){
		seq_count[i] = 0;
		sum_gc3[i] = 0;
		for(j=0; j<unique_codons; ++j){
			codon_bins[i][j] = 0;
		}
	}

	for(i=0; i<unique_codons; ++i)
		codons[i] = 0;

	seq_length = getSequence(seqFile, seq);

	//Parse the fasta file
	while(seq_length != -1){
		used_seq++;

		//Check the quality of the sequence
		if(seq_length % 3 != 0)
			bad_seq++;

		//If sequence is good use it
		else{
			int codon_count = seq_length/3;
			int gc3 = 0;
			double per_gc3;
			int gc_bin;
			int cdn_index;

			//Get the gc3 content
			for(i=0; i<codon_count; ++i){
				memcpy(cur_cdn,&seq[i*3],3);		
				if(cur_cdn[2] == 'G' || cur_cdn[2] == 'C')
					gc3++;
			}
			
			per_gc3 = 100*(gc3/(double)codon_count);
			gc_bin = (int)(per_gc3/binsize);
			
			seq_count[gc_bin]++;
			sum_gc3[gc_bin] += per_gc3;

			//count occurences of each codon in sequence
			for(i=0; i<codon_count; ++i){
				memcpy(cur_cdn, &seq[i*3],3);
				cdn_index = get_codon_index(cur_cdn);
				codon_bins[gc_bin][cdn_index]++;
			}
			
		}
		seq_length = getSequence(seqFile, seq);
	}

	//Done with the fasta sequence file...close it up
	fclose(seqFile);

	
	//Combine the bins to form partially overlappig bins
	double ave_gc3[bins-1];
	int ave_seq[bins-1];
	int final_bins = 0;
	double sum_tot;
	int seq_tot;

	for(i=0; i<bins-1; ++i){
		seq_tot = seq_count[i] + seq_count[i+1];
		sum_tot = sum_gc3[i] + sum_gc3[i+1];

		//Take average %GC3 for bin unless there are no sequences in bin
		if(seq_tot > 0){
			ave_gc3[i] = sum_tot/seq_tot;
			ave_seq[i] = seq_tot;
		}else{
			ave_gc3[i] = 0;
			ave_seq[i] = 0;
		}
	}
	
	//Amino acid bin counter
	int aa_counts[bins-1][unique_aa];
	for(i=0; i<bins-1; ++i)
		for(j=0; j<unique_aa; ++j)
			aa_counts[i][j] = 0;

	//Codon Data
	int codon_data[unique_codons][bins];
	for(i=0; i<unique_codons; ++i)
		for(j=0; j<bins; ++j)
			codon_data[i][j] = 0;

	int codon1;
	int codon2;
	int codon_tot;
	int aa_index;
	char aa;

	for(i=0; i<unique_codons; ++i){

		for(j=0; j<bins-1; ++j){
			codon1 = codon_bins[j][i];
			codon2 = codon_bins[j+1][i];
			codon_tot = codon1 + codon2;
	
			//Match the codon to the AA
			aa = aa_codon_map[i];
			aa_index = get_aa_index(aa);

			aa_counts[j][aa_index] += codon_tot;	
			codon_data[i][j] = codon_tot;	
		}	
	}		
	
	//Array for multipliers from mulFile	
	double mult_rows[unique_codons][unique_codons];
	char mult_read[10000];
	double tmp;
	int d,len;
	char *result = NULL;
	char delims[] = "\t";

	//Read the first two lines
	fscanf(mulFile,"%[^\n]%*c", mult_read);
	//fprintf(stderr,"%s\n", mult_read);
	

	fscanf(mulFile,"%[^\n]%*c", mult_read);
	//fprintf(stderr,"%s\n", mult_read);
	
	for(i=0; i<unique_codons; ++i){
	
		fscanf(mulFile,"%s", mult_read);
		//fprintf(stderr,"%s\n", mult_read);
		fscanf(mulFile,"%s", mult_read);
		//fprintf(stderr,"%s\n", mult_read);

		for(j=0; j<unique_codons; ++j){
			fscanf(mulFile,"%s",mult_read);
			mult_rows[i][j] = atof(mult_read);
			//fprintf(stderr, "%f\t", mult_rows[i][j]);
		} 	
		
		//fprintf(stderr,"\n");
		
	}
	fprintf(stdout, "Parsed the mult file\n");	
	fclose(mulFile);

	//Start doing the output
	fprintf(outFile, "Bin\tSequences\tGC3 Ave\tAA\tCodons");	

	int aa_last;
	fprintf(stdout, "Starting the output\n");	
	for(i=0; i<bins-1; ++i){
		aa_last = 100;
		for(j=0; j<unique_codons; ++j){
			aa = aa_codon_map[j];
			aa_index = get_aa_index(aa);
			
			int codon_sum = codon_data[j][i];
			int aa_total = aa_counts[i][aa_index];
			double codon_per;
			
			if(aa_total > 0){
				codon_per = 100*(codon_sum/(double)aa_total);
			}

			while(aa_last != aa_index){
				double ave = ave_gc3[i];
				int seqc = ave_seq[i];
				
				fprintf(outFile, "\n%d\t", i); 
				fprintf(outFile, "%d\t", ave_seq[i]);
				fprintf(outFile, "%d\t", (int)ave_gc3[i]);
				fprintf(outFile, "%c", aa);
				
				aa_last = aa_index;
			}
	
			for(k=0; k<unique_codons; ++k){
				double multiplier = mult_rows[j][k];
				int newnum;
				int printed = 0;
				
				if(aa_total > 0){
					newnum = (int)(multiplier * codon_per + 0.5);			
				}
				else{
					newnum = 0;
				}

				if(k == 0)
					fprintf(outFile,"\t");
				else
					fprintf(outFile," ");

				if(aa_total > 0 && newnum > 0){
						while(printed != newnum){
							fprintf(outFile,"%s",codon_types[j]);
							++printed;
						}
				} else{ 
					fprintf(outFile, "NNN");

				}
			}
	
		}
	}	

	fprintf(stdout, "Printing codon totals for each bin\n");
	//Print codon totals for each bin, w/o multipliers, to be used for picking 3' codons
	for(i=0; i<bins-1; ++i){
		aa_last = 100;
		double codon_sum;
		
		for(j=0; j<unique_codons; ++j){
			codon_sum = codon_data[j][i];
			aa = aa_codon_map[j];
			aa_index = get_aa_index(aa);
	
			int aa_total = aa_counts[i][aa_index];
			double codon_per;
			int newnum;

			if(aa_total > 0)
				codon_per = 100*(codon_sum/aa_total);
		

			while(aa_last != aa_index){
				if(i == 0 && j ==0) {fprintf(outFile, "\n");}
				
				fprintf(outFile,"\n%d\t", i);
				fprintf(outFile,"TOTALS\t%d", ave_seq[i]);
				fprintf(outFile,"\t%f", ave_gc3[i]);
				fprintf(outFile,"\t%c", aa);
		
				aa_last = aa_index;
			}
	
			if(aa_total > 0){
				newnum = (int)codon_per + .5;
				int printed = 0;
				while(printed != newnum){
					fprintf(outFile, "%s", codon_types[j]);
					printed++;
				}
			}else{
				fprintf(outFile, " NNN"); //if no AA
	
			}
		}
	}
	
	fclose(outFile);
	
	fprintf(stdout, "Total Sequences: %d\n", used_seq);
	fprintf(stdout, "Bad Sequences: %d\n", bad_seq);
	
}
